import re
import string
from functools import reduce

import pandas as pd
from openpyxl.formula.translate import Translator
from openpyxl.formula import Tokenizer
from openpyxl import load_workbook
from settings import compare_column,contrast_compare_column


def handle_ref(cell_range, idx, amount,how,before=False):
    '''
     # (1)如果是单个单元格引用如A5，如果5大于等于idx，则5+amount，否则不变；
     # (2)如果是范围引用如A6:A9，如果6大于等于idx，则6+amount，否则不变；如果9大于等于idx，则9+amount，否则不变；
    This function handles cell reference and returns the new cell reference after inserting rows.
    :param cell_range:引用单元格范围
    :param idx:插入的行或列的索引
    :param amount:插入的行数或列数
    :param how:column or row
    :param before:是否在idx之前的行或列
    :return:
    '''

    if ":" in cell_range:
        # 处理范围引用
        start, end = cell_range.split(":")
        start_col, start_row = get_cell_column_row(start)
        end_col, end_row = get_cell_column_row(end)
        if how == "row":
            if before:
                if start_row > idx - 1:
                    start_row += amount
                if end_row > idx - 1:
                    end_row += amount
            else:
                if start_row >= idx-1:
                    start_row += amount
                if end_row >= idx-1:
                    end_row += amount
        elif how == "column":
            if before:
                if start_col > idx-1:
                    start_col += amount
                if end_col > idx-1:
                    end_col += amount
            else:
                if start_col >= idx-1:
                    start_col += amount
                if end_col >= idx-1:
                    end_col += amount
        else:
            raise ValueError("Invalid how parameter.")
        new_cell_range = "{}{}:{}{}".format(contrast_compare_column[start_col], start_row, contrast_compare_column[end_col], end_row)
    else:
        # 处理单个单元格引用
        col, row = get_cell_column_row(cell_range)
        if how == "row":
            if before:
                if row > idx-1:
                    row += amount
            else:
                if row >= idx-1:
                    row += amount
        elif how == "column":
            if before:
                if col > idx-1:
                    col += amount
            else:
                if col >= idx-1:
                    col += amount
        else:
            raise ValueError("Invalid how parameter.")
        new_cell_range = "{}{}".format(contrast_compare_column[col], row)
    return new_cell_range


def insert_rows_with_formula(wb,s_name,idx,amount):
    '''
    This function inserts rows with formulas.
    :param wb:
    :param s_name:sheet name
    :param idx:
    :param amount:
    :return:
    '''
    ws = wb[s_name]
    # 先插入空行
    ws.insert_rows(idx, amount)
    # 找到倒数第一行
    last_row = list(ws.rows)[idx-2]
    # 循环处理最后一行的每一个单元格，如果有公式，则替换公式
    for i,cell in enumerate(last_row):
        if cell.data_type == "f":
            formula = cell.value
            # 对下面所有的单元格进行公式替换
            for j in range(idx,idx+amount):
                dest_cell = ws.cell(row=j, column=i+1)
                translated = Translator(formula, cell.coordinate).translate_formula(dest_cell.coordinate)
                ws[dest_cell.coordinate] = translated
    # 改变所有后面行和前面行的公式，原单元格中的公式需要加上对应的行数增加amount
    change_rows = list(range(1,idx)) + list(range(idx+amount,ws.max_row+1))
    for i in change_rows:
        for j in range(1,ws.max_column+1):
            if ws.cell(row=i, column=j).data_type == "f":
                formula = ws.cell(row=i, column=j).value
                tokenizer = Tokenizer(formula)
                new_formula = "="
                for item in tokenizer.items:
                    if item.subtype == "RANGE":
                        # 处理RANGE类型，
                        # 如果是本sheet的引用，
                        # (1)如果是单个单元格引用如A5，如果5大于等于idx，则5+amount，否则不变；
                        # (2)如果是范围引用如A6:A9，如果6大于等于idx，则6+amount，否则不变；如果9大于等于idx，则9+amount，否则不变；
                        # 如果是其他sheet的引用，则不加amount，如Sheet1!A1不变
                        if "!" in item.value:
                            # 处理Sheet1!A1这种引用
                            sheet_name, cell_range = item.value.split("!")
                            if sheet_name == s_name:
                                new_cell_range = handle_ref(cell_range, idx, amount,"row",i<idx)
                                new_formula += "{}!{}".format(sheet_name, new_cell_range)
                            else:
                                new_formula += item.value
                        else:
                            new_cell_range = handle_ref(item.value, idx, amount,"row",i<idx)
                            new_formula += new_cell_range
                    else:
                        new_formula += item.value
                ws.cell(row=i, column=j, value=new_formula)


def insert_columns_with_formula(wb,s_name,idx,amount):
    '''
    This function inserts columns with formulas.
    :param wb:
    :param s_name:sheet name
    :param idx:
    :param amount:
    :return:
    '''
    ws = wb[s_name]
    # 先插入空列
    ws.insert_cols(idx, amount)
    # 找到倒数第一列
    last_col = list(ws.columns)[idx-2]
    # 循环处理最后一列的每一个单元格，如果有公式，则替换公式
    for i,cell in enumerate(last_col):
        if cell.data_type == "f":
            formula = cell.value
            # 对下面所有的单元格进行公式替换
            for j in range(idx,idx+amount):
                dest_cell = ws.cell(row=i+1, column=j)
                translated = Translator(formula, cell.coordinate).translate_formula(dest_cell.coordinate)
                ws[dest_cell.coordinate] = translated
    # 改变所有后面列的公式，原单元格中的公式需要加上对应的列数增加amount
    change_cols = list(range(1,idx)) + list(range(idx+amount,ws.max_column+1))
    for i in change_cols:
        for j in range(1,ws.max_row+1):
            if ws.cell(row=j, column=i).data_type == "f":
                formula = ws.cell(row=j, column=i).value
                tokenizer = Tokenizer(formula)
                new_formula = "="
                for item in tokenizer.items:
                    if item.subtype == "RANGE":
                        # 处理RANGE类型，
                        # 如果是本sheet的引用，
                        # (1)如果是单个单元格引用如A5，如果5大于等于idx，则5+amount，否则不变；
                        # (2)如果是范围引用如A6:A9，如果6大于等于idx，则6+amount，否则不变；如果9大于等于idx，则9+amount，否则不变；
                        # 如果是其他sheet的引用，则不加amount，如Sheet1!A1不变
                        if "!" in item.value:
                            # 处理Sheet1!A1这种引用
                            sheet_name, cell_range = item.value.split("!")
                            if sheet_name == s_name:
                                new_cell_range = handle_ref(cell_range, idx, amount,"column",i<idx)
                                new_formula += "{}!{}".format(sheet_name, new_cell_range)
                            else:
                                new_formula += item.value
                        else:
                            new_cell_range = handle_ref(item.value, idx, amount,"column",i<idx)
                            new_formula += new_cell_range
                    else:
                        new_formula += item.value
                ws.cell(row=j, column=i, value=new_formula)


def get_cell_column_row(cell):
    '''
    This function returns the column and row number of a cell in a spreadsheet.
    :param cell: The cell whose column and row number is to be returned.
    :return: (int, int): The column and row number of the cell.
    '''

    # 拆分单元格坐标,文本为列，数字为行，使用正则表达式
    col, row = re.findall(r'\D+|\d+', cell)
    return compare_column[col], int(row)


def add_data(ws,data,start_cell,how="column",formula=False):
    '''
    This function adds data to a cell in a spreadsheet.
    :param ws: ws (Worksheet): The spreadsheet where the data should be added.
    :param data:  The data to be added to the cell.
    :param start_cell: (str): The cell where the data should be added.
    :param how: The method to be used to add the data.values can be "column" or "row" or "matrix".
    :param formula: Whether to add the data as a formula.
    :return:
    '''

    # 获取数据长度
    data_length = len(data)
    # 拆分单元格坐标,文本为列，数字为行，使用正则表达式
    col,row = get_cell_column_row(start_cell)
    # 循环添加数据
    if how == "column":
        for i in range(data_length):
            # 如果cell中有公式，则不添加数据
            if formula and ws.cell(row=row+i, column=col).data_type == "f":
                continue
            ws.cell(row=row+i, column=col, value=data[i])
    elif how == "row":
        for i in range(data_length):
            # 如果cell中有公式，则不添加数据
            if formula and ws.cell(row=row + i, column=col).data_type == "f":
                continue
            ws.cell(row=row, column=col+i, value=data[i])
    elif how == "matrix":
        # 采用列优先
        for i in range(data_length):
            for j in range(len(data[i])):
                # 如果cell中有公式，则不添加数据
                if formula and ws.cell(row=row+j, column=col+i).data_type == "f":
                    continue
                ws.cell(row=row+j, column=col+i, value=data[i][j])
    else:    # 未知的添加方式
        print("Unknown adding method.")


def add_mutable_data(df_from,df_to,from_col,to_col,index_col,wb,sheet_name,start_cell,how="column",formula=True):
    '''
    步骤：
    TODO：如果index_col列中有重复值或者有缺失值，则需要处理，
    This function adds mutable data to a cell in a spreadsheet.
    :param df_from: The original dataframe.
    :param df_to: The updated dataframe.
    :param from_col: The column name of the original dataframe.it can be a list of column names.
    :param to_col: The column name of the updated dataframe.it can be a list of column names.
    :param index_col: The index column name.
    :param wb: The workbook where the data should be added.
    :param sheet_name: The sheet name where the data should be added.
    :param start_cell: (str): The cell where the data should be added.
    :param how: The method to be used to add the data.values can be "column" or "row" or "matrix".
    :param formula: Whether to add the data as a formula.
    :return:
    '''
    # 第一步：获取所有的索引值
    # 获取df_to的index_col列数据
    index_col_num = df_to.columns.tolist().index(index_col) + 1
    if isinstance(to_col,list):
        df_to = df_to[[index_col]+to_col]
    else:
        df_to = df_to[[index_col,to_col]]
    to_index_list = list(df_to[index_col])
    to_index_set = set(to_index_list)

    # 合并数据
    if isinstance(from_col,list):
        df_from = df_from[[index_col]+from_col]
        print(type(df_from[index_col]))
        print(df_from.loc[:,index_col].dtype)
        for i in from_col:
            print(i)
            print(df_from.loc[:,i].dtype)
        conditons = [df_from[i].abs() > 0.00 if df_from[i].dtype == 'float64' else df_from[i].notnull() for i in from_col]
        # 使用or连接多个条件
        df_from = df_from[reduce(lambda x,y:x|y,conditons)]
    else:
        df_from = df_from[[index_col,from_col]]
        if df_from[from_col].dtype == 'float64':
            df_from = df_from[df_from[from_col].abs() > 0.00]

    from_index_list = list(df_from[index_col])
    from_index_set = set(from_index_list)
    # 找出to_index_set中不存在的index
    diff_index_set = from_index_set - to_index_set
    diff_index = list(diff_index_set)
    if len(to_index_list) ==0:
        whole_index = diff_index
    else:
        whole_index = to_index_list[:len(to_index_list)-1] + diff_index + [to_index_list[-1]]

    # 第二步：获取索引对应的数据值
    df_combine = pd.merge(df_from, df_to, on=index_col, how='outer')
    df_combine = df_combine.copy()
    if isinstance(from_col,list):
        whole_value = []
        for i in from_col:
            if df_combine[i].dtype == 'float64':
                df_combine[i] = df_combine[i].fillna(0)
            # 将'key'列设置为索引，然后转换为字典
            dict_from_df = df_combine.set_index(index_col)[i].to_dict()
            # 在ws末尾添加空行，并添加diff_index的数据
            column_value = [dict_from_df[i] for i in whole_index]
            whole_value.append(column_value)
    else:
        if df_combine[from_col].dtype == 'float64':
            df_combine[from_col] = df_combine[from_col].fillna(0)
        # 将'key'列设置为索引，然后转换为字典
        dict_from_df = df_combine.set_index(index_col)[from_col].to_dict()
        # 在ws末尾添加空行，并添加diff_index的数据
        whole_value = [dict_from_df[i] for i in whole_index]

    # 在ws末尾添加空行，并添加diff的数据
    ws = wb[sheet_name]
    max_row = ws.max_row
    if len(diff_index) > 0:
        insert_rows_with_formula(wb, sheet_name, max_row, len(diff_index))
        # 在ws的index_col列中添加diff的数据
        index_cell = "{}{}".format(contrast_compare_column[index_col_num], max_row)
        add_data(ws, diff_index, index_cell, "column", formula)
    # 在ws的index_col列中添加diff的数据
    add_data(ws, whole_value, start_cell, how, formula)




if __name__ == '__main__':
    wb = load_workbook(r'D:\audit\audit_report\project\model-义乌城投\model - 副本.xlsx')
    ws = wb['应收利息分类']
    # insert_rows_with_formula(wb, '应收利息分类', 6, 2)
    # insert_columns_with_formula(wb, "应收利息分类",3, 2)
    df_from  = pd.read_excel(r'D:\audit\audit_report\project\model-义乌城投\义乌市城投供应链管理有限公司_TB - 副本.xlsx',sheet_name='应收利息分类')
    df_to = pd.read_excel(r'D:\audit\audit_report\project\model-义乌城投\model - 副本.xlsx',sheet_name='应收利息分类')
    add_mutable_data(df_from,df_to,'期末余额','期初余额','项目',wb,'应收利息分类','C2',how='column',formula=True),
    wb.save(r'D:\audit\audit_report\project\model-义乌城投\model - 副本.xlsx')
    wb.close()
    # df_to = pd.DataFrame({'index':[1,2,3,4,5,11,12,13,14,15],'to_col':[1,2,3,4,5,11,12,13,14,15]})
    # df_to = df_to.astype({'index':str,'to_col':int})
    # df_from = pd.DataFrame({'index':[1,2,3,4,5,6,7,8,9,10],'from_col':[1,2,3,4,5,6,7,8,9,10]})
    # df_from = df_from.astype({'index':str,'from_col':int})
    # df_combine = pd.merge(df_to, df_from,on='index', how='outer',sort=False)
    # print(df_combine)
    # df_combine_new = df_combine[df_combine['to_col'].isnull()]
    # print(df_combine_new)

